#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
自然语言处理内容识别分析技术研发计划实现
基于BERT模型的命名实体识别(NER)系统
"""

import json
import os
import numpy as np
import logging
import torch
from typing import List, Dict, Any
from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import pipeline

# 配置日志
# 确保日志目录存在
log_dir = "logs"
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

# 为这个模块创建独立的logger
logger = logging.getLogger('main')
logger.setLevel(logging.INFO)

# 清除可能存在的旧处理器
for handler in logger.handlers[:]:
    logger.removeHandler(handler)

# 创建文件处理器
file_handler = logging.FileHandler(os.path.join(log_dir, "app.log"), encoding='utf-8')
file_handler.setLevel(logging.INFO)

# 创建控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# 创建格式化器
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)
console_handler.setFormatter(formatter)

# 添加处理器到logger
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# 防止日志传播到根logger
logger.propagate = False

# 导入增强版NER处理器
from advanced_ner import AdvancedNERProcessor
# 导入微调模型NER处理器
from fine_tuned_ner import FineTunedNERProcessor
# 导入数据库管理器
from database_manager import DatabaseManager
# 导入BERT微调器
from bert_finetuner import BERTFineTuner

app = Flask(__name__)

class NERProcessor:
    """
    基于BERT的命名实体识别处理器
    """

    def __init__(self, model_path="/home/aresen/1project/2python/hub/models/chinese-bert-wwm-ext"):
        """
        初始化NER处理器，加载预训练模型
        """
        # 使用本地中文BERT模型进行NER任务
        self.model_path = model_path
        self.use_bert = False
        # 自定义实体文件路径
        self.custom_entities_file = "custom_entities.json"

        # 检查本地模型路径是否存在
        if os.path.exists(self.model_path):
            try:
                print(f"从本地路径加载BERT模型: {self.model_path}")
                self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, local_files_only=True)
                self.model = AutoModelForTokenClassification.from_pretrained(self.model_path, local_files_only=True)
                
                # 检查CUDA是否可用并移动模型
                if torch.cuda.is_available():
                    self.model = self.model.to('cuda')
                    print("BERT模型已移动到GPU")
                
                self.ner_pipeline = pipeline(
                    "ner",
                    model=self.model,
                    tokenizer=self.tokenizer,
                    aggregation_strategy="simple"
                )
                self.use_bert = True
                print(f"成功加载本地BERT模型: {self.model_path}")
            except Exception as e:
                print(f"警告: 无法从本地路径加载BERT模型 {self.model_path}: {e}")
        else:
            print(f"警告: 本地模型路径不存在 {self.model_path}")
            # 尝试使用在线模型作为备选
            try:
                self.model_name = "chinese-bert-wwm-ext"
                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
                self.model = AutoModelForTokenClassification.from_pretrained(self.model_name)
                
                # 检查CUDA是否可用并移动模型
                if torch.cuda.is_available():
                    self.model = self.model.to('cuda')
                    print("BERT模型已移动到GPU")
                
                self.ner_pipeline = pipeline(
                    "ner",
                    model=self.model,
                    tokenizer=self.tokenizer,
                    aggregation_strategy="simple"
                )
                self.use_bert = True
                print("成功加载在线BERT模型: chinese-bert-wwm-ext")
            except Exception as e:
                print(f"警告: 无法加载在线BERT模型: {e}")

        # 动态实体词典，用于添加自定义实体
        self.custom_entities = {}
        # 加载持久化的自定义实体
        self.load_custom_entities()

    def add_custom_entity(self, entity_key: str, entity_value: str, domain: str = "general"):
        """
        动态添加命名实体识别规则

        Args:
            entity_key: 实体键
            entity_value: 实体值
            domain: 领域，默认为通用领域
        """
        if domain not in self.custom_entities:
            self.custom_entities[domain] = {}

        if entity_key not in self.custom_entities[domain]:
            self.custom_entities[domain][entity_key] = []

        if entity_value not in self.custom_entities[domain][entity_key]:
            self.custom_entities[domain][entity_key].append(entity_value)
            
        # 每次添加新实体后保存到文件
        self.save_custom_entities()

    def save_custom_entities(self):
        """
        将自定义实体保存到文件
        """
        try:
            with open(self.custom_entities_file, 'w', encoding='utf-8') as f:
                json.dump(self.custom_entities, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"保存自定义实体失败: {e}")

    def load_custom_entities(self):
        """
        从文件加载自定义实体
        """
        try:
            if os.path.exists(self.custom_entities_file):
                with open(self.custom_entities_file, 'r', encoding='utf-8') as f:
                    self.custom_entities = json.load(f)
                print(f"成功加载自定义实体，共 {len(self.custom_entities)} 个领域")
        except Exception as e:
            print(f"加载自定义实体失败: {e}")

    def extract_entities(self, text: str, domain: str = "general") -> List[Dict[str, Any]]:
        """
        从文本中提取命名实体

        Args:
            text: 输入文本
            domain: 领域

        Returns:
            实体列表
        """
        entities = []

        # 使用BERT模型进行基础NER识别（如果可用）
        if self.use_bert:
            try:
                bert_entities = self.ner_pipeline(text)
                for entity in bert_entities:
                    # 确保置信度是Python原生float类型，而不是numpy.float32
                    confidence = float(entity["score"]) if isinstance(entity["score"], (np.float32, np.float64)) else entity["score"]
                    entities.append({
                        "nerKey": entity["entity_group"],
                        "nerValue": entity["word"],
                        "confidence": confidence
                    })
            except Exception as e:
                print(f"BERT实体识别出错: {e}")

        # 处理自定义实体
        custom_entities = self._extract_custom_entities(text, domain)

        # 合并结果
        entities.extend(custom_entities)

        return entities

    def _extract_custom_entities(self, text: str, domain: str) -> List[Dict[str, Any]]:
        """
        从文本中提取自定义实体

        Args:
            text: 输入文本
            domain: 领域

        Returns:
            自定义实体列表
        """
        entities = []

        # 检查通用领域和指定领域的自定义实体
        domains_to_check = ["general"]
        if domain != "general":
            domains_to_check.append(domain)

        for check_domain in domains_to_check:
            if check_domain in self.custom_entities:
                for entity_key, entity_values in self.custom_entities[check_domain].items():
                    for entity_value in entity_values:
                        if entity_value in text:
                            entities.append({
                                "nerKey": entity_key,
                                "nerValue": entity_value,
                                "confidence": 1.0,
                                "source": "custom"
                            })

        return entities

def convert_numpy_types(obj):
    """
    递归转换numpy类型为Python原生类型，确保JSON序列化正常工作
    """
    if isinstance(obj, dict):
        return {key: convert_numpy_types(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [convert_numpy_types(item) for item in obj]
    elif isinstance(obj, (np.float32, np.float64)):
        return float(obj)
    elif isinstance(obj, (np.int32, np.int64)):
        return int(obj)
    else:
        return obj

# 初始化处理器
ner_processor = NERProcessor()
advanced_ner_processor = AdvancedNERProcessor()

# 尝试初始化微调模型处理器（如果存在）
fine_tuned_ner_processor = None
try:
    # 检查微调模型目录是否存在且包含必要的文件
    fine_tuned_model_path = "./fine_tuned_model"
    if os.path.exists(fine_tuned_model_path):
        # 检查必要的模型文件
        required_files = ["config.json", "pytorch_model.bin", "label_mapping.json"]
        files_exist = all(os.path.exists(os.path.join(fine_tuned_model_path, f)) for f in required_files)

        if files_exist:
            fine_tuned_ner_processor = FineTunedNERProcessor(fine_tuned_model_path)
            print("成功加载微调模型处理器")
        else:
            print(f"微调模型目录 {fine_tuned_model_path} 缺少必要文件，跳过加载")
    else:
        print(f"微调模型目录 {fine_tuned_model_path} 不存在，跳过加载")
except Exception as e:
    print(f"加载微调模型处理器失败: {e}")
    import traceback
    traceback.print_exc()

# 初始化数据库管理器
db_manager = DatabaseManager()

@app.route('/nlp/analysisText', methods=['POST'])
def analyze_text():
    """
    命名实体识别接口
    """
    try:
        # 获取请求数据
        data = request.get_json()

        system_code = data.get('systemCode')
        text = data.get('text')
        domain = data.get('domain', 'general')
        method = data.get('method', 'bert')  # bert, advanced, 或 fine_tuned

        # 简单的系统码验证
        if not system_code or system_code != "NLP_SYSTEM_001":
            return jsonify({
                "error": "Invalid system code"
            }), 401

        if not text:
            return jsonify({
                "error": "Text is required"
            }), 400

        # 根据方法选择处理器
        if method == "advanced":
            entities = advanced_ner_processor.extract_entities(text, domain)
        elif method == "fine_tuned":
            if fine_tuned_ner_processor:
                entities = fine_tuned_ner_processor.extract_entities(text)
            else:
                return jsonify({
                    "error": "Fine-tuned model not available. Please fine-tune a model first."
                }), 400
        else:
            entities = ner_processor.extract_entities(text, domain)

        # 转换numpy类型以确保JSON序列化正常工作
        entities = convert_numpy_types(entities)

        return jsonify({
            "ners": entities
        })

    except Exception as e:
        logger.error(f"实体识别接口出错: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return jsonify({
            "error": str(e)
        }), 500

@app.route('/nlp/addCustomEntity', methods=['POST'])
def add_custom_entity():
    """
    添加自定义实体接口
    """
    try:
        data = request.get_json()

        entity_key = data.get('entityKey')
        entity_value = data.get('entityValue')
        domain = data.get('domain', 'general')

        if not entity_key or not entity_value:
            return jsonify({
                "error": "entityKey and entityValue are required"
            }), 400

        ner_processor.add_custom_entity(entity_key, entity_value, domain)

        return jsonify({
            "message": "Custom entity added successfully"
        })

    except Exception as e:
        logger.error(f"添加自定义实体接口出错: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return jsonify({
            "error": str(e)
        }), 500

@app.route('/nlp/finetune', methods=['POST'])
def finetune_model():
    """
    微调模型接口
    """
    try:
        logger.info("开始处理模型微调请求")
        # 检查系统码
        data = request.get_json()
        system_code = data.get('systemCode')
        if not system_code or system_code != "NLP_SYSTEM_001":
            logger.warning("无效的系统码")
            return jsonify({
                "error": "Invalid system code"
            }), 401

        logger.info("连接数据库获取训练数据")
        # 连接数据库
        if not db_manager.connect():
            logger.error("数据库连接失败")
            return jsonify({
                "error": "Failed to connect to database"
            }), 500

        # 获取训练数据
        train_data = db_manager.get_bert_train_data()
        db_manager.disconnect()
        
        logger.info(f"获取到 {len(train_data)} 条训练数据")

        if not train_data:
            logger.warning("没有可用的训练数据")
            return jsonify({
                "error": "No training data available"
            }), 400

        # 微调模型
        logger.info("开始微调模型")
        finetuner = BERTFineTuner()
        model_path = finetuner.fine_tune(train_data)

        # 验证模型文件是否已保存
        logger.info("验证模型文件是否已保存")
        required_files = ["config.json", "pytorch_model.bin", "label_mapping.json"]
        missing_files = []
        for file in required_files:
            file_path = os.path.join(model_path, file)
            if not os.path.exists(file_path):
                missing_files.append(file)
        
        if missing_files:
            logger.error(f"模型文件缺失: {missing_files}")
            return jsonify({
                "error": f"模型文件缺失: {missing_files}"
            }), 500

        # 重新加载微调模型处理器
        logger.info("重新加载微调模型处理器")
        global fine_tuned_ner_processor
        fine_tuned_ner_processor = FineTunedNERProcessor(model_path)

        logger.info("模型微调成功完成")
        return jsonify({
            "message": "Model fine-tuned successfully",
            "model_path": model_path
        })

    except Exception as e:
        logger.error(f"模型微调接口出错: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return jsonify({
            "error": str(e)
        }), 500

@app.route('/nlp/databaseInfo', methods=['GET'])
def get_database_info():
    """
    获取数据库信息接口
    """
    try:
        # 检查系统码
        system_code = request.args.get('systemCode')
        if not system_code or system_code != "NLP_SYSTEM_001":
            return jsonify({
                "error": "Invalid system code"
            }), 401

        # 连接数据库
        if not db_manager.connect():
            return jsonify({
                "error": "Failed to connect to database"
            }), 500

        # 获取数据库信息
        fields = db_manager.get_bert_fields()
        field_details = db_manager.get_bert_field_details()
        nouns = db_manager.get_bert_nouns()

        db_manager.disconnect()

        return jsonify({
            "fields": fields,
            "field_details": field_details,
            "nouns": nouns
        })

    except Exception as e:
        return jsonify({
            "error": str(e)
        }), 500

if __name__ == '__main__':
    # 添加一些示例自定义实体
    ner_processor.add_custom_entity("组织", "承德钒钛")
    ner_processor.add_custom_entity("工种", "钳工")
    ner_processor.add_custom_entity("单位", "河钢集团")
    ner_processor.add_custom_entity("压力单位", "MPa")
    ner_processor.add_custom_entity("温度单位", "℃")
    ner_processor.add_custom_entity("证书", "操作证")
    ner_processor.add_custom_entity("方案", "行动工作方案")

    app.run(host='0.0.0.0', port=5000, debug=True)